{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "067b892d",
   "metadata": {},
   "source": [
    "# 🗣️ Humanize DB - Query your SQL Database with Natural Language\n",
    "\n",
    "## 📋 Project Overview\n",
    "\n",
    "**Humanize DB** is an intelligent SQL database agent that allows users to interact with SQL databases using natural human language. Instead of writing complex SQL queries, simply ask questions in plain English and get accurate, formatted answers.\n",
    "\n",
    "This project leverages **LangGraph** to create a sophisticated multi-agent workflow that:\n",
    "- Understands natural language questions\n",
    "- Explores database schemas automatically\n",
    "- Generates and validates SQL queries\n",
    "- Executes queries safely\n",
    "- Returns human-readable answers\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "\n",
    "## 📝 Project Structure\n",
    "\n",
    "```\n",
    "humandb.ipynb\n",
    "├── Setup & Configuration\n",
    "│   ├── Import dependencies\n",
    "│   ├── Initialize LLM (GPT-4)\n",
    "│   └── Download sample database\n",
    "├── Database Tools\n",
    "│   ├── SQL query executor\n",
    "│   ├── Schema inspector\n",
    "│   └── Table lister\n",
    "├── Agent Nodes\n",
    "│   ├── list_tables() - Discover tables\n",
    "│   ├── call_get_schema() - Get schema\n",
    "│   ├── generate_query() - Create SQL\n",
    "│   ├── check_query() - Validate SQL\n",
    "│   └── format responses\n",
    "├── Workflow Definition\n",
    "│   └── LangGraph StateGraph\n",
    "└── User Interfaces\n",
    "    ├── Programmatic interface\n",
    "    └── Gradio Chat UI\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c08875b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import init_chat_model\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from langchain_community.agent_toolkits import SQLDatabaseToolkit\n",
    "from typing import Literal\n",
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "from langgraph.graph import END, START, MessagesState, StateGraph\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from IPython.display import Image, display\n",
    "from langchain_core.runnables.graph import CurveStyle, MermaidDrawMethod, NodeStyles\n",
    "import requests\n",
    "import gradio as gr\n",
    "\n",
    "llm = init_chat_model(\"openai:gpt-4.1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ca4b689",
   "metadata": {},
   "source": [
    "## Configure the database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a5a1b9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db\"\n",
    "DB_NAME = \"sample.db\"\n",
    "\n",
    "response = requests.get(url)\n",
    "\n",
    "if response.status_code == 200:\n",
    "    with open(DB_NAME, \"wb\") as file:\n",
    "        file.write(response.content)\n",
    "    print(f\"file downloaded and saved as {DB_NAME}\")\n",
    "else:\n",
    "    print(f\"Failed to download File. Error : {response.status_code}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e0a0f13",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(f\"sqlite:///{DB_NAME}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32ba2a48",
   "metadata": {},
   "outputs": [],
   "source": [
    "toolkit = SQLDatabaseToolkit(db=db, llm=llm)\n",
    "tools = toolkit.get_tools()\n",
    "\n",
    "for tool in tools:\n",
    "    print(f\"{tool.name}: {tool.description}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bae9b573",
   "metadata": {},
   "source": [
    "### Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47e44902",
   "metadata": {},
   "source": [
    "I am using some of Langchain's tools specifically built for querying SQL, and displaying the schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "701b7b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_schema_tool = next(tool for tool in tools if tool.name == \"sql_db_schema\")\n",
    "run_query_tool = next(tool for tool in tools if tool.name == \"sql_db_query\")\n",
    "list_tables_tool = next(tool for tool in tools if tool.name == \"sql_db_list_tables\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f70159",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_schema_node = ToolNode([get_schema_tool], name=\"get_schema\")\n",
    "run_query_node = ToolNode([run_query_tool], name=\"run_query\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cde35a4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_tables(state: MessagesState):\n",
    "    tool_call = {\n",
    "        \"name\": \"sql_db_list_tables\",\n",
    "        \"args\": {},\n",
    "        \"id\": \"abc123\",\n",
    "        \"type\": \"tool_call\",\n",
    "    }\n",
    "\n",
    "    tool_call_message = AIMessage(content=\"\", tool_calls=[tool_call])\n",
    "    tool_message = list_tables_tool.invoke(tool_call)\n",
    "    response = AIMessage(f\"Available tables: {tool_message.content}\")\n",
    "    return {\"messages\": [tool_call_message, tool_message, response]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "611763e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def call_get_schema(state:MessagesState):\n",
    "    llm_with_tools = llm.bind_tools([get_schema_tool], tool_choice=\"any\")\n",
    "    response = llm_with_tools.invoke(state[\"messages\"])\n",
    "    return {\"messages\":[response]}\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7068d1f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_query_system_prompt = \"\"\"\n",
    "You are an agent designed to interact with a SQL database.\n",
    "Given an input question, create a syntactically correct {dialect} query to run,\n",
    "then look at the results of the query and return the answer. Unless the user\n",
    "specifies a specific number of examples they wish to obtain, always limit your\n",
    "query to at most {top_k} results.\n",
    "\n",
    "You can order the results by a relevant column to return the most interesting\n",
    "examples in the database. Never query for all the columns from a specific table,\n",
    "only ask for the relevant columns given the question.\n",
    "\n",
    "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\n",
    "\"\"\".format(\n",
    "    dialect=db.dialect,\n",
    "    top_k=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22457569",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Node to generate the SQL QUERY\n",
    "\n",
    "\n",
    "def generate_query(state: MessagesState):\n",
    "    system_message = {\"role\": \"system\", \"content\": generate_query_system_prompt}\n",
    "\n",
    "    llm_with_tools = llm.bind_tools([run_query_tool])\n",
    "    response = llm_with_tools.invoke([system_message] + state[\"messages\"])\n",
    "\n",
    "    return {\"messages\": [response]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdd7fcf2",
   "metadata": {},
   "source": [
    "### Query Evaluator Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "964a7051",
   "metadata": {},
   "outputs": [],
   "source": [
    "validate_query_system_prompt = \"\"\"\n",
    "You are a SQL expert with a strong attention to detail.\n",
    "Double check the {dialect} query for common mistakes, including:\n",
    "- Using NOT IN with NULL values\n",
    "- Using UNION when UNION ALL should have been used\n",
    "- Using BETWEEN for exclusive ranges\n",
    "- Data type mismatch in predicates\n",
    "- Properly quoting identifiers\n",
    "- Using the correct number of arguments for functions\n",
    "- Casting to the correct data type\n",
    "- Using the proper columns for joins\n",
    "\n",
    "If there are any of the above mistakes, rewrite the query. If there are no mistakes,\n",
    "just reproduce the original query.\n",
    "\n",
    "You will call the appropriate tool to execute the query after running this check.\n",
    "\"\"\".format(dialect=db.dialect)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6740bd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_query(state: MessagesState):\n",
    "    system_message = {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": validate_query_system_prompt,\n",
    "    }\n",
    "\n",
    "    # Generate an artificial user message to check\n",
    "    tool_call = state[\"messages\"][-1].tool_calls[0]\n",
    "    user_message = {\"role\": \"user\", \"content\": tool_call[\"args\"][\"query\"]}\n",
    "    llm_with_tools = llm.bind_tools([run_query_tool], tool_choice=\"any\")\n",
    "    response = llm_with_tools.invoke([system_message, user_message])\n",
    "    response.id = state[\"messages\"][-1].id\n",
    "\n",
    "    return {\"messages\": [response]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5087c047",
   "metadata": {},
   "outputs": [],
   "source": [
    "def should_continue(state: MessagesState) -> Literal[END, \"check_query\"]:\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    if not last_message.tool_calls:\n",
    "        return END\n",
    "    else:\n",
    "        return \"check_query\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78c02422",
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = StateGraph(MessagesState)\n",
    "builder.add_node(list_tables)\n",
    "builder.add_node(call_get_schema)\n",
    "builder.add_node(get_schema_node, \"get_schema\")\n",
    "builder.add_node(generate_query)\n",
    "builder.add_node(check_query)\n",
    "builder.add_node(run_query_node, \"run_query\")\n",
    "\n",
    "builder.add_edge(START, \"list_tables\")\n",
    "builder.add_edge(\"list_tables\", \"call_get_schema\")\n",
    "builder.add_edge(\"call_get_schema\", \"get_schema\")\n",
    "builder.add_edge(\"get_schema\", \"generate_query\")\n",
    "builder.add_conditional_edges(\n",
    "    \"generate_query\",\n",
    "    should_continue,\n",
    ")\n",
    "builder.add_edge(\"check_query\", \"run_query\")\n",
    "builder.add_edge(\"run_query\", \"generate_query\")\n",
    "\n",
    "agent = builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fe41b10",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Image(agent.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c97e733",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat_with_agent(message, history):\n",
    "    \"\"\"\n",
    "    Process user message and return agent's response\n",
    "    \"\"\"\n",
    "    response = \"\"\n",
    "    \n",
    "    # Stream the agent's response\n",
    "    for step in agent.stream(\n",
    "        {\"messages\": [{\"role\": \"user\", \"content\": message}]},\n",
    "        stream_mode=\"values\",\n",
    "    ):\n",
    "        # Get the last message content\n",
    "        last_message = step[\"messages\"][-1]\n",
    "        if hasattr(last_message, 'content'):\n",
    "            response = last_message.content\n",
    "    \n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f884e7c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_tables():\n",
    "    tool_call = {\n",
    "        \"name\": \"sql_db_list_tables\",\n",
    "        \"args\": {},\n",
    "        \"id\": \"abc123\",\n",
    "        \"type\": \"tool_call\",\n",
    "    }\n",
    "    tables = list_tables_tool(tool_call)\n",
    "    return tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db071275",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks(theme=gr.themes.Soft()) as demo:\n",
    "    gr.Markdown(\"#  Humaanize DB - Query your SQL Database with Human Language\")\n",
    "\n",
    "    with gr.Row():\n",
    "        # Left column: Chat interface\n",
    "        with gr.Column(scale=3):\n",
    "            gr.Markdown(\"Ask questions about music tracks, genres, and get insights!\")\n",
    "\n",
    "            chatbot = gr.Chatbot(\n",
    "                height=500,\n",
    "                bubble_full_width=False,\n",
    "            )\n",
    "\n",
    "            with gr.Row():\n",
    "                msg = gr.Textbox(\n",
    "                    placeholder=\"Ask about genres, tracks, or statistics...\",\n",
    "                    show_label=False,\n",
    "                    scale=4,\n",
    "                )\n",
    "                submit = gr.Button(\"Send\", scale=1, variant=\"primary\")\n",
    "\n",
    "            with gr.Row():\n",
    "                clear = gr.Button(\"Clear Chat\")\n",
    "\n",
    "            gr.Examples(\n",
    "                examples=[\n",
    "                    \"Which genre on average has the longest tracks?\",\n",
    "                    \"What are the top 5 most popular genres?\",\n",
    "                    \"How many tracks are in each genre?\",\n",
    "                ],\n",
    "                inputs=msg,\n",
    "            )\n",
    "\n",
    "        # Right column: Tables panel\n",
    "        with gr.Column(scale=1):\n",
    "            gr.Markdown(\"### 📊 Available Tables\")\n",
    "            tables_display = gr.Textbox(\n",
    "                label=\"Database Tables\",\n",
    "                value=display_tables(),\n",
    "                lines=10,\n",
    "                interactive=False,\n",
    "                max_lines=15,\n",
    "            )\n",
    "            refresh_btn = gr.Button(\"🔄 Refresh Tables\", size=\"sm\")\n",
    "            refresh_btn.click(display_tables, outputs=tables_display)\n",
    "\n",
    "    # Chat functionality\n",
    "    def user(user_message, history):\n",
    "        return \"\", history + [[user_message, None]]\n",
    "\n",
    "    def bot(history):\n",
    "        user_message = history[-1][0]\n",
    "        bot_response = chat_with_agent(user_message, history)\n",
    "        history[-1][1] = bot_response\n",
    "        return history\n",
    "\n",
    "    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(\n",
    "        bot, chatbot, chatbot\n",
    "    )\n",
    "    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(\n",
    "        bot, chatbot, chatbot\n",
    "    )\n",
    "    clear.click(lambda: None, None, chatbot, queue=False)\n",
    "\n",
    "demo.launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
